import torch
from config.device_config import get_device
from src.model.layers.knn import KNN

device = get_device()


def triangle_generator_loss(original_nodes, original_barycenters, selected_triangles, selected_triangles_probabilities,
                            closure_loss_weight=0.1):  # 增加封闭性损失权重参数，默认0.1
    """
    计算三角形生成器的损失，包括形状生成的损失和封闭性损失。
    """
    b = original_barycenters
    b_hat = selected_triangles.mean(dim=1)
    p_b_hat = selected_triangles_probabilities

    d_f_S_Ss = _torch_d_f_S_Ss(p_b_hat, b_hat, b)
    d_r_S_Ss = _torch_d_r_S_Ss(original_nodes, selected_triangles, selected_triangles_probabilities, b_hat)

    # 计算封闭性损失，并加权
    closure_loss = _mesh_closure_loss(selected_triangles) * closure_loss_weight  # 根据权重调整封闭性损失

    total_loss = d_f_S_Ss + d_r_S_Ss + closure_loss
    return d_f_S_Ss, d_r_S_Ss, closure_loss


def _torch_d_f_S_Ss(p_b_hat, b_hat, b):
    """
    计算三角形的最小距离，并基于概率加权求和。
    """
    distances = torch.cdist(b_hat, b)
    distances_filtered = torch.where(distances != 0, distances, torch.inf)  # 避免零距离影响
    min_b = distances_filtered.min(dim=1).values
    final_term = torch.sum(p_b_hat * min_b)
    return final_term


def _sample_points_in_triangles(triangles, num_points_per_triangle=50):
    """
    在三角形内部均匀采样点。每个三角形中采样的点的数量由num_points_per_triangle指定。
    """
    N = triangles.shape[0]
    r1 = torch.sqrt(torch.rand(N, num_points_per_triangle, device=torch.device(device)))
    r2 = torch.rand(N, num_points_per_triangle, device=torch.device(device))

    a = 1 - r1
    b = r1 * (1 - r2)
    c = r1 * r2

    a = a.unsqueeze(2)
    b = b.unsqueeze(2)
    c = c.unsqueeze(2)

    sampled_points = (a * triangles[:, 0].unsqueeze(1) +
                      b * triangles[:, 1].unsqueeze(1) +
                      c * triangles[:, 2].unsqueeze(1))

    return sampled_points


def _torch_d_r_S_Ss(original_nodes, generated_triangles, probability_generated_triangles,
                    barycenters_generated_triangles):
    """
    计算生成网格与原始网格之间的损失。
    """
    k = 20
    sampled_points_generated_triangles = _sample_points_in_triangles(generated_triangles)

    distances = torch.cdist(original_nodes, sampled_points_generated_triangles)
    min_distances = distances.min(dim=1).values
    min_mean_distances = min_distances.mean(dim=1)
    first_term = probability_generated_triangles * min_mean_distances

    indices_knn = KNN(k=k, batch_size=1000)(barycenters_generated_triangles).int()

    x_tk = sampled_points_generated_triangles[indices_knn[:, 1:]]
    x_y = sampled_points_generated_triangles[indices_knn[:, 0]].unsqueeze(1)
    knn_distances = torch.cdist(x_tk, x_y)
    knn_distances_mean = knn_distances.mean(dim=3)
    probabilities_tk = probability_generated_triangles[indices_knn[:, 1:]]
    dist_times_probability = knn_distances_mean * probabilities_tk.unsqueeze(2)
    sum_over_k = dist_times_probability.sum(dim=1)
    mean_points_per_triangle_again = sum_over_k.mean(dim=1)
    normed = mean_points_per_triangle_again / k
    second_term = normed * (1 - probability_generated_triangles)

    loss_d_r_S_Ss = torch.sum(first_term + second_term)

    return loss_d_r_S_Ss


def _mesh_closure_loss(triangles):
    """
    计算网格的封闭性损失，确保生成的三角形网格是封闭的。
    """
    edges = _get_edges_from_triangles(triangles)
    edge_counts = {}

    for edge in edges:
        if isinstance(edge[0], torch.Tensor):
            edge_list = [item.tolist() for item in edge]
        else:
            edge_list = edge

        edge_list_flat = []
        for item in edge_list:
            if isinstance(item, torch.Tensor):
                edge_list_flat.append(float(item.item()))
            elif isinstance(item, list):
                edge_list_flat.extend([float(i) for i in item])
            else:
                edge_list_flat.append(float(item))

        edge_tuple = tuple(sorted(edge_list_flat))
        edge_counts[edge_tuple] = edge_counts.get(edge_tuple, 0) + 1

    non_closing_edges = sum(1 for count in edge_counts.values() if count != 2)

    return torch.tensor(non_closing_edges, dtype=torch.float32)


def _get_edges_from_triangles(triangles):
    """
    从三角形中提取边，每个三角形有 3 条边。
    """
    edges = []
    for triangle in triangles:
        edges.append([triangle[0], triangle[1]])
        edges.append([triangle[1], triangle[2]])
        edges.append([triangle[2], triangle[0]])
    return edges
